"""
简化的RandomForest模型运行脚本
避免编码问题
"""

import os
import sys
import json
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import KFold
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import joblib
from datetime import datetime

class SimpleRandomForestRunner:
    def __init__(self):
        """初始化运行器"""
        # 获取项目根目录
        script_dir = os.path.dirname(os.path.abspath(__file__))
        model_dir = os.path.dirname(script_dir)
        self.project_root = os.path.dirname(model_dir)

        # 创建输出目录
        self.output_dir = os.path.join(script_dir, 'outputs')
        os.makedirs(os.path.join(self.output_dir, 'models'), exist_ok=True)
        os.makedirs(os.path.join(self.output_dir, 'reports'), exist_ok=True)
        os.makedirs(os.path.join(self.output_dir, 'submissions'), exist_ok=True)

        self.model = None
        self.feature_names = None

    def load_data(self):
        """加载数据"""
        print("[INFO] Loading data...")

        try:
            train_path = os.path.join(self.project_root, 'feature/train_feature.csv')
            test_path = os.path.join(self.project_root, 'feature/test_feature.csv')

            train_data = pd.read_csv(train_path)
            test_data = pd.read_csv(test_path)

            # 准备数据
            X_train = train_data.drop(['SalePrice', 'Id'], axis=1, errors='ignore')
            y_train = np.log1p(train_data['SalePrice'])  # 对数变换
            X_test = test_data.drop(['Id'], axis=1, errors='ignore')
            test_ids = test_data['Id'] if 'Id' in test_data.columns else None

            self.feature_names = list(X_train.columns)

            print(f"[OK] Data loaded successfully")
            print(f"   Train shape: {X_train.shape}")
            print(f"   Test shape: {X_test.shape}")
            print(f"   Features: {len(self.feature_names)}")

            return X_train, y_train, X_test, test_ids

        except Exception as e:
            print(f"[ERROR] Failed to load data: {e}")
            return None, None, None, None

    def train_with_cv(self, X_train, y_train):
        """使用交叉验证训练模型"""
        print("[INFO] Training RandomForest with cross-validation...")

        # 模型参数
        model_params = {
            'n_estimators': 200,
            'max_depth': None,
            'min_samples_split': 2,
            'min_samples_leaf': 1,
            'max_features': 'sqrt',
            'bootstrap': True,
            'random_state': 42,
            'n_jobs': -1
        }

        # 交叉验证
        kf = KFold(n_splits=5, shuffle=True, random_state=42)

        rmse_scores = []
        mae_scores = []
        r2_scores = []

        start_time = datetime.now()

        for fold, (train_idx, val_idx) in enumerate(kf.split(X_train), 1):
            print(f"Training fold {fold}...")

            X_train_fold = X_train.iloc[train_idx]
            X_val_fold = X_train.iloc[val_idx]
            y_train_fold = y_train.iloc[train_idx]
            y_val_fold = y_train.iloc[val_idx]

            # 训练模型
            model = RandomForestRegressor(**model_params)
            model.fit(X_train_fold, y_train_fold)

            # 预测
            y_pred = model.predict(X_val_fold)

            # 计算指标
            rmse = np.sqrt(mean_squared_error(y_val_fold, y_pred))
            mae = mean_absolute_error(y_val_fold, y_pred)
            r2 = r2_score(y_val_fold, y_pred)

            rmse_scores.append(rmse)
            mae_scores.append(mae)
            r2_scores.append(r2)

            print(f"Fold {fold}: RMSE={rmse:.4f}, MAE={mae:.4f}, R2={r2:.4f}")

        # 计算平均指标
        avg_rmse = np.mean(rmse_scores)
        std_rmse = np.std(rmse_scores)
        avg_mae = np.mean(mae_scores)
        std_mae = np.std(mae_scores)
        avg_r2 = np.mean(r2_scores)
        std_r2 = np.std(r2_scores)

        end_time = datetime.now()
        training_time = str(end_time - start_time).split('.')[0]

        print(f"\nCross-validation results:")
        print(f"RMSE: {avg_rmse:.4f} +/- {std_rmse:.4f}")
        print(f"MAE: {avg_mae:.4f} +/- {std_mae:.4f}")
        print(f"R2: {avg_r2:.4f} +/- {std_r2:.4f}")
        print(f"Training time: {training_time}")

        # 在全部数据上训练最终模型
        print("Training final model on all data...")
        self.model = RandomForestRegressor(**model_params)
        self.model.fit(X_train, y_train)

        # 获取特征重要性
        feature_importances = pd.DataFrame({
            'feature': self.feature_names,
            'importance': self.model.feature_importances_
        }).sort_values('importance', ascending=False)

        print("\nTop 10 important features:")
        print(feature_importances.head(10))

        # 生成验证报告
        self.generate_validation_report(avg_rmse, std_rmse, rmse_scores,
                                     avg_mae, std_mae, mae_scores,
                                     avg_r2, std_r2, r2_scores,
                                     training_time, model_params)

        return True

    def generate_validation_report(self, avg_rmse, std_rmse, rmse_scores,
                                 avg_mae, std_mae, mae_scores,
                                 avg_r2, std_r2, r2_scores,
                                 training_time, model_params):
        """生成验证报告"""
        print("[INFO] Generating validation report...")

        report = {
            "model_info": {
                "model_name": "RandomForest Baseline Model",
                "version": "1.0.0",
                "framework": "sklearn",
                "created_at": datetime.now().isoformat()
            },
            "data_info": {
                "train_size": len(self.feature_names),
                "n_features": len(self.feature_names),
                "target": "SalePrice",
                "cv_method": "KFold",
                "n_splits": 5,
                "shuffle": True,
                "random_state": 42
            },
            "metrics": {
                "RMSE": {
                    "mean": float(avg_rmse),
                    "std": float(std_rmse),
                    "folds": [float(score) for score in rmse_scores]
                },
                "MAE": {
                    "mean": float(avg_mae),
                    "std": float(std_mae),
                    "folds": [float(score) for score in mae_scores]
                },
                "R2": {
                    "mean": float(avg_r2),
                    "std": float(std_r2),
                    "folds": [float(score) for score in r2_scores]
                }
            },
            "training_summary": {
                "training_time": training_time,
                "cross_val_accuracy": float(avg_r2),
                "notes": "Baseline model with RandomForest algorithm."
            },
            "environment": {
                "python": "3.12.6",
                "os": "Windows",
                "machine": "AMD64"
            },
            "model_args": model_params
        }

        # 保存报告
        report_path = os.path.join(self.output_dir, 'reports', 'validation_report_randomforest.json')
        with open(report_path, 'w', encoding='utf-8') as f:
            json.dump(report, f, indent=2)

        print(f"[OK] Validation report saved: {report_path}")

    def generate_predictions(self, X_test, test_ids):
        """生成预测"""
        print("[INFO] Generating predictions...")

        try:
            # 预测
            y_pred_log = self.model.predict(X_test)
            y_pred = np.expm1(y_pred_log)  # 还原对数变换

            # 创建提交文件
            submission = pd.DataFrame({
                'Id': test_ids,
                'SalePrice': y_pred
            })

            submission_path = os.path.join(self.output_dir, 'submissions', 'submission_randomforest.csv')
            submission.to_csv(submission_path, index=False)

            print(f"[OK] Submission file saved: {submission_path}")
            print(f"   Min price: ${y_pred.min():,.2f}")
            print(f"   Max price: ${y_pred.max():,.2f}")
            print(f"   Avg price: ${y_pred.mean():,.2f}")

            return True

        except Exception as e:
            print(f"[ERROR] Failed to generate predictions: {e}")
            return False

    def save_model(self):
        """保存模型"""
        print("[INFO] Saving model...")

        try:
            # 保存模型
            model_path = os.path.join(self.output_dir, 'models', 'randomforest_model.pkl')
            joblib.dump(self.model, model_path)

            # 保存特征名称
            feature_path = os.path.join(self.output_dir, 'models', 'feature_names.pkl')
            joblib.dump(self.feature_names, feature_path)

            # 保存特征重要性
            feature_importance_df = pd.DataFrame({
                'feature': self.feature_names,
                'importance': self.model.feature_importances_
            }).sort_values('importance', ascending=False)

            importance_path = os.path.join(self.output_dir, 'models', 'feature_importance.csv')
            feature_importance_df.to_csv(importance_path, index=False)

            print(f"[OK] Model saved: {model_path}")
            print(f"[OK] Features saved: {feature_path}")
            print(f"[OK] Feature importance saved: {importance_path}")

            return True

        except Exception as e:
            print(f"[ERROR] Failed to save model: {e}")
            return False

    def run(self):
        """运行完整流程"""
        print("[INFO] Starting RandomForest training...")
        print("=" * 50)

        # 加载数据
        X_train, y_train, X_test, test_ids = self.load_data()
        if X_train is None:
            return False

        # 训练模型
        if not self.train_with_cv(X_train, y_train):
            return False

        # 保存模型
        if not self.save_model():
            return False

        # 生成预测
        if not self.generate_predictions(X_test, test_ids):
            return False

        print("\n" + "=" * 50)
        print("[SUCCESS] RandomForest training completed!")
        print("Output files:")
        print("  - Model: outputs/models/randomforest_model.pkl")
        print("  - Submission: outputs/submissions/submission_randomforest.csv")
        print("  - Report: outputs/reports/validation_report_randomforest.json")

        return True

if __name__ == "__main__":
    runner = SimpleRandomForestRunner()
    success = runner.run()

    if not success:
        print("[ERROR] Training failed!")
        sys.exit(1)